#
# fashion_minst_t01.py
#
###############################################################################
#
import time
import torch
import torchvision
import fashionMnist as fmnist
import torch.utils.data as data
import matplotlib.pyplot as plt
import torch.nn.functional as F

FASHION_MNIST_ROOT = fmnist.ROOT

# 配置超参数
LR = 0.0001
EPOCHS = 20
WORKERS = 0   # 对于Windows用户，这里应设置为0，否则会出现多线程错误
BATCH_SIZE = 256
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 获取数据集
train_data = fmnist.get_fashion_mnist(FASHION_MNIST_ROOT, train=True)
test_data = fmnist.get_fashion_mnist(FASHION_MNIST_ROOT, train=False)
# print(type(train_data))

train_data_size = len(train_data)
test_data_size = len(test_data)

print(f"训练数据集的长度为 {train_data_size}")          # 60000
print(f"测试数据集的长度为 {test_data_size}")           # 10000
print(f"训练数据集的形状维 {train_data.data.shape}")    # [60000, 28, 28]

# 数据集形状
feature, lable = train_data[0]
print(feature.shape, feature.dtype)
print(lable)

# DataLoader
train_loader = data.DataLoader(
    train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = data.DataLoader(
    test_data, batch_size=BATCH_SIZE, shuffle=True)

#


def my_test_00():
    """测试图像"""
    train_loader_iter = iter(train_loader)
    X, y = next(train_loader_iter)
    img = torch.Tensor(X[0]).reshape(28, 28)
    # torch.Tensor(X[0]).shape, img.shape
    print(torch.Tensor(y[0]).item(), fmnist.get_fashion_mnist_kind_name(y[0]))
    plt.imshow(img, cmap=plt.cm.gray)
    plt.show()


def my_test_01():
    """读取小批量数据"""
    start = time.time()
    for X, _ in train_loader:
        print(X.shape)
        continue
    print('%.2f sec' % (time.time() - start))


# my_test_00()
